"""
动态时间规整(DTW)算法

实现了DTW算法用于时间序列对齐和相似度计算，包括：
- 基础DTW算法
- 能熵比预处理
- 相似度修正
- 公共子串优化
"""

import numpy as np
from typing import Tuple, Optional, List
from loguru import logger
from scipy.fft import fft
from scipy.signal import windows


class DTWAlgorithm:
    """动态时间规整算法"""
    
    def __init__(self, 
                 window_size: int = 10,
                 step_pattern: str = "symmetric2",
                 distance_metric: str = "euclidean",
                 normalization: bool = True,
                 frame_length: int = 256,
                 frame_shift: int = 128,
                 window_type: str = "hamming",
                 fft_size: int = 512,
                 energy_constant: float = 1.0):
        """
        初始化DTW算法
        
        Args:
            window_size: DTW窗口大小
            step_pattern: 步进模式
            distance_metric: 距离度量
            normalization: 是否归一化
            frame_length: 帧长度
            frame_shift: 帧移
            window_type: 窗函数类型
            fft_size: FFT大小
            energy_constant: 能量常数
        """
        self.window_size = window_size
        self.step_pattern = step_pattern
        self.distance_metric = distance_metric
        self.normalization = normalization
        self.frame_length = frame_length
        self.frame_shift = frame_shift
        self.window_type = window_type
        self.fft_size = fft_size
        self.energy_constant = energy_constant
        
        logger.info("初始化DTW算法")
    
    def compute_dtw_distance(self, 
                           sequence_x: np.ndarray, 
                           sequence_y: np.ndarray) -> Tuple[float, np.ndarray]:
        """
        计算DTW距离
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            
        Returns:
            (DTW距离, 规整路径)
        """
        # 能熵比预处理
        x_processed = self._energy_entropy_preprocessing(sequence_x)
        y_processed = self._energy_entropy_preprocessing(sequence_y)
        
        # 计算DTW距离
        distance, path = self._dtw_core(x_processed, y_processed)
        
        # 相似度修正
        corrected_distance = self._similarity_correction(x_processed, y_processed, distance)
        
        return corrected_distance, path
    
    def _energy_entropy_preprocessing(self, signal_data: np.ndarray) -> np.ndarray:
        """
        能熵比预处理
        
        Args:
            signal_data: 输入信号
            
        Returns:
            能熵比序列
        """
        # 去除直流分量和归一化
        signal_centered = signal_data - np.mean(signal_data)
        max_amplitude = np.max(np.abs(signal_centered))
        if max_amplitude > 0:
            signal_normalized = signal_centered / max_amplitude
        else:
            signal_normalized = signal_centered
        
        # 分帧
        num_frames = (len(signal_normalized) - self.frame_length) // self.frame_shift + 1
        energy_entropy_ratios = np.zeros(num_frames)
        
        # 选择窗函数
        if self.window_type == "rectangular":
            window = windows.boxcar(self.frame_length)
        elif self.window_type == "hamming":
            window = windows.hamming(self.frame_length)
        elif self.window_type == "hanning":
            window = windows.hann(self.frame_length)
        else:
            window = windows.hamming(self.frame_length)
        
        for i in range(num_frames):
            start = i * self.frame_shift
            end = start + self.frame_length
            frame = signal_normalized[start:end]
            
            # 加窗
            frame_windowed = frame * window
            
            # FFT变换
            spectrum = fft(frame_windowed, self.fft_size)
            magnitude_spectrum = np.abs(spectrum[:self.fft_size//2])
            
            # 计算能量
            energy = np.log(1 + np.sum(magnitude_spectrum) / self.energy_constant)
            
            # 计算谱熵
            magnitude_sum = np.sum(magnitude_spectrum)
            if magnitude_sum > 0:
                probabilities = magnitude_spectrum / magnitude_sum
                # 避免log(0)
                probabilities = np.maximum(probabilities, 1e-10)
                entropy = -np.sum(probabilities * np.log(probabilities))
            else:
                entropy = 0
            
            # 计算能熵比
            if entropy > 0:
                energy_entropy_ratios[i] = np.sqrt(1 + np.abs(energy / entropy))
            else:
                energy_entropy_ratios[i] = energy
        
        return energy_entropy_ratios
    
    def _dtw_core(self, 
                  sequence_x: np.ndarray, 
                  sequence_y: np.ndarray) -> Tuple[float, np.ndarray]:
        """
        DTW核心算法
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            
        Returns:
            (DTW距离, 规整路径)
        """
        len_x, len_y = len(sequence_x), len(sequence_y)
        
        # 初始化距离矩阵
        distance_matrix = np.full((len_x + 1, len_y + 1), np.inf)
        distance_matrix[0, 0] = 0
        
        # 填充距离矩阵
        for i in range(1, len_x + 1):
            for j in range(1, len_y + 1):
                # 计算当前点距离
                current_distance = self._calculate_point_distance(
                    sequence_x[i-1], sequence_y[j-1]
                )
                
                # 寻找最小路径
                min_prev = min(
                    distance_matrix[i-1, j],      # 插入
                    distance_matrix[i, j-1],      # 删除
                    distance_matrix[i-1, j-1]     # 匹配
                )
                
                distance_matrix[i, j] = current_distance + min_prev
        
        # 回溯路径
        path = self._backtrack_path(distance_matrix)
        
        # 计算最终距离
        final_distance = distance_matrix[len_x, len_y]
        
        # 归一化
        if self.normalization:
            path_length = len(path)
            if path_length > 0:
                final_distance = final_distance / path_length
        
        return final_distance, path
    
    def _calculate_point_distance(self, x: float, y: float) -> float:
        """
        计算两点间距离
        
        Args:
            x: 点x的值
            y: 点y的值
            
        Returns:
            距离
        """
        if self.distance_metric == "euclidean":
            return (x - y) ** 2
        elif self.distance_metric == "manhattan":
            return abs(x - y)
        elif self.distance_metric == "cosine":
            # 对于标量值，余弦距离退化为简单的差值
            return 1 - np.cos(abs(x - y))
        else:
            return (x - y) ** 2
    
    def _backtrack_path(self, distance_matrix: np.ndarray) -> np.ndarray:
        """
        回溯DTW路径
        
        Args:
            distance_matrix: 距离矩阵
            
        Returns:
            规整路径
        """
        len_x, len_y = distance_matrix.shape
        i, j = len_x - 1, len_y - 1
        
        path = []
        while i > 0 and j > 0:
            path.append((i-1, j-1))
            
            # 寻找前一个最小距离点
            if distance_matrix[i-1, j] <= distance_matrix[i, j-1] and distance_matrix[i-1, j] <= distance_matrix[i-1, j-1]:
                i -= 1
            elif distance_matrix[i, j-1] <= distance_matrix[i-1, j-1]:
                j -= 1
            else:
                i -= 1
                j -= 1
        
        # 处理边界情况
        while i > 0:
            path.append((i-1, 0))
            i -= 1
        while j > 0:
            path.append((0, j-1))
            j -= 1
        
        path.reverse()
        return np.array(path)
    
    def _similarity_correction(self, 
                             sequence_x: np.ndarray, 
                             sequence_y: np.ndarray, 
                             original_distance: float) -> float:
        """
        相似度修正
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            original_distance: 原始距离
            
        Returns:
            修正后的距离
        """
        # 计算公共子串
        common_substring_length = self._calculate_common_substring(sequence_x, sequence_y)
        
        # 计算惩罚系数
        min_length = min(len(sequence_x), len(sequence_y))
        if min_length > 0:
            penalty_coefficient = 1 - common_substring_length / min_length
        else:
            penalty_coefficient = 1
        
        # 应用惩罚系数
        corrected_distance = penalty_coefficient * original_distance
        
        return corrected_distance
    
    def _calculate_common_substring(self, 
                                  sequence_x: np.ndarray, 
                                  sequence_y: np.ndarray) -> int:
        """
        计算公共子串长度
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            
        Returns:
            公共子串长度
        """
        len_x, len_y = len(sequence_x), len(sequence_y)
        
        # 计算标准差
        std_x = np.std(sequence_x)
        std_y = np.std(sequence_y)
        max_std = max(std_x, std_y)
        
        # 设置偏移容忍范围
        tolerance = max_std * 0.5
        
        # 动态规划计算最长公共子串
        dp = np.zeros((len_x + 1, len_y + 1), dtype=int)
        
        for i in range(1, len_x + 1):
            for j in range(1, len_y + 1):
                if abs(sequence_x[i-1] - sequence_y[j-1]) < tolerance:
                    dp[i, j] = dp[i-1, j-1] + 1
                else:
                    dp[i, j] = max(dp[i-1, j], dp[i, j-1])
        
        return dp[len_x, len_y]
    
    def compute_similarity(self, 
                          sequence_x: np.ndarray, 
                          sequence_y: np.ndarray) -> float:
        """
        计算相似度
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            
        Returns:
            相似度 (0-1之间，1表示完全相同)
        """
        distance, _ = self.compute_dtw_distance(sequence_x, sequence_y)
        
        # 将距离转换为相似度
        similarity = 1 / (1 + distance)
        
        return similarity
    
    def find_best_match(self, 
                       query_sequence: np.ndarray, 
                       reference_sequences: List[np.ndarray]) -> Tuple[int, float]:
        """
        在参考序列中找到最佳匹配
        
        Args:
            query_sequence: 查询序列
            reference_sequences: 参考序列列表
            
        Returns:
            (最佳匹配索引, 相似度)
        """
        best_index = -1
        best_similarity = -1
        
        for i, ref_sequence in enumerate(reference_sequences):
            similarity = self.compute_similarity(query_sequence, ref_sequence)
            
            if similarity > best_similarity:
                best_similarity = similarity
                best_index = i
        
        return best_index, best_similarity
    
    def plot_dtw_path(self, 
                     sequence_x: np.ndarray, 
                     sequence_y: np.ndarray,
                     path: np.ndarray = None):
        """
        绘制DTW路径
        
        Args:
            sequence_x: 序列X
            sequence_y: 序列Y
            path: DTW路径
        """
        try:
            import matplotlib.pyplot as plt
            
            if path is None:
                _, path = self.compute_dtw_distance(sequence_x, sequence_y)
            
            # 预处理序列
            x_processed = self._energy_entropy_preprocessing(sequence_x)
            y_processed = self._energy_entropy_preprocessing(sequence_y)
            
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
            
            # 绘制原始序列
            ax1.plot(x_processed, 'b-', label='序列X', linewidth=2)
            ax1.plot(y_processed, 'r-', label='序列Y', linewidth=2)
            ax1.set_title('原始序列对比')
            ax1.set_xlabel('时间')
            ax1.set_ylabel('能熵比')
            ax1.legend()
            ax1.grid(True, alpha=0.3)
            
            # 绘制DTW路径
            if len(path) > 0:
                ax2.plot(path[:, 0], path[:, 1], 'g-', linewidth=2, label='DTW路径')
                ax2.plot([0, len(x_processed)-1], [0, len(y_processed)-1], 'k--', alpha=0.5, label='对角线')
            
            ax2.set_title('DTW规整路径')
            ax2.set_xlabel('序列X索引')
            ax2.set_ylabel('序列Y索引')
            ax2.legend()
            ax2.grid(True, alpha=0.3)
            ax2.set_aspect('equal')
            
            plt.tight_layout()
            plt.show()
            
        except ImportError:
            logger.warning("matplotlib未安装，无法绘制DTW路径")
    
    def get_algorithm_parameters(self) -> dict:
        """
        获取算法参数
        
        Returns:
            参数字典
        """
        return {
            'window_size': self.window_size,
            'step_pattern': self.step_pattern,
            'distance_metric': self.distance_metric,
            'normalization': self.normalization,
            'frame_length': self.frame_length,
            'frame_shift': self.frame_shift,
            'window_type': self.window_type,
            'fft_size': self.fft_size,
            'energy_constant': self.energy_constant
        }
    
    def set_algorithm_parameters(self, **kwargs):
        """
        设置算法参数
        
        Args:
            **kwargs: 要设置的参数
        """
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
                logger.info(f"设置{key}={value}")
            else:
                logger.warning(f"未知参数: {key}")
    
    def __str__(self) -> str:
        return (f"DTWAlgorithm(window_size={self.window_size}, "
                f"distance_metric={self.distance_metric}, normalization={self.normalization})")
    
    def __repr__(self) -> str:
        return self.__str__() 